import torch
import torch.nn as nn
import numpy as np


class TestNet(nn.Module):
    def __init__(self):
        super(TestNet, self).__init__()
        
    def forward(self, predict, target):
        mse_loss = torch.mean(torch.pow((predict - target), 2))
        return [mse_loss]


if __name__ == "__main__":
    x = np.array([[0,1],[2,3]]).astype(np.float32)
    y = np.array([[0,2],[-1,3]]).astype(np.float32)

    x = torch.Tensor(x)
    y = torch.Tensor(y)

    net = TestNet()
    loss = net(x, y)
    print(loss)

